/*
This file is part of asproxy.

asproxy is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

asproxy is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with asproxy.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "proxyhelper.h"
#include "log.h"
#include <arpa/inet.h>

#define STX 15
#define ETX 14
#define DEL 28

SSL_CTX*         g_psslcontext = NULL;
pthread_mutex_t* g_lockarray = NULL;
char* (*g_phandle_incoming)(char* bytes, int* length) = NULL;
char* (*g_phandle_outgoing)(char* bytes, int* length) = NULL;

void setsockettimeout(int sock, int timeout)
{
	struct timeval socktimeout;
	socktimeout.tv_sec = timeout;
	socktimeout.tv_usec = 0;
	setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (char *)&socktimeout, sizeof(socktimeout));
	setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, (char *)&socktimeout, sizeof(socktimeout));
}

int connectserver()
{
	int sockfd = 0;
	struct sockaddr_in serv_addr;

	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd < 0) 
	{
		log_error("Could not create a valid server side socket");
		
		if (errno > 0)
			log_error(strerror(errno));
		return -1;
	}
	
	bzero((char *) &serv_addr, sizeof(serv_addr));
	serv_addr.sin_family = AF_INET;
	inet_pton(AF_INET, g_proxyconfig.DestinationIP, &(serv_addr.sin_addr));
	serv_addr.sin_port = htons(g_proxyconfig.DestinationPort);

	if (connect(sockfd,(struct sockaddr *) &serv_addr,sizeof(serv_addr)) < 0)
	{
		log_error("Could not connect to primary server.");
		
		if (errno > 0)
			log_error(strerror(errno));
		
		// Try connecting to the failover if provided
		if (strcmp(g_proxyconfig.FailoverIP, NONE_DEFAULT) != 0)
		{			
			struct sockaddr_in failover_addr;
			bzero((char *) &failover_addr, sizeof(failover_addr));
			failover_addr.sin_family = AF_INET;
			inet_pton(AF_INET, g_proxyconfig.FailoverIP, &(failover_addr.sin_addr));
			failover_addr.sin_port = htons(g_proxyconfig.DestinationPort);		
			
			if (connect(sockfd,(struct sockaddr *) &failover_addr,sizeof(failover_addr)) < 0)
			{
				log_error("Could not connect to failover server.");
				if (errno > 0)
					log_error(strerror(errno));
					
				if (sockfd > 0)
					close(sockfd);
				return -1;
			}
			else 
				log_debug("Connected to failover server.");
		}
		else 
		{
			if (sockfd > 0)
			{
				shutdown(sockfd, SHUT_RD);
				close(sockfd);
			}
			return -1;
		}
	}
	
	log_info("Server socket connection completed.");
	return sockfd;
}

SSL_CTX* initserverctx(void)
{	
	if (g_psslcontext != NULL)
		return g_psslcontext;
	
	// Initialize the library
	SSL_library_init();
	
    OpenSSL_add_all_algorithms();  /* load & register all cryptos, etc. */
    SSL_load_error_strings();   /* load all error messages */
    g_psslcontext = SSL_CTX_new(SSLv23_server_method());   /* create new context from method */
    if (g_psslcontext == NULL)
    {
        log_error("SSL_CTX_new failed");
		return NULL;
    }
	
	// Load trusted CAs
	SSL_CTX_set_verify_depth(g_psslcontext, 1);
	
	// load the certificates
	loadcertificates(g_psslcontext);
		
	return g_psslcontext;
}

void loadcertificates(SSL_CTX* ctx)
{
	/* set the local certificate from CertFile */
    if (SSL_CTX_use_certificate_chain_file(ctx, g_proxyconfig.SSLCertFile) <= 0)
        log_error("Use Cert File failed");
		
    /* set the private key from KeyFile (may be the same as CertFile) */
    if (SSL_CTX_use_PrivateKey_file(ctx, g_proxyconfig.SSLCertFile, SSL_FILETYPE_PEM) <= 0)
        log_error("Use Private Key failed");
	
	if (SSL_CTX_load_verify_locations(ctx, g_proxyconfig.SSLCertFile, g_proxyconfig.SSLCAPath) <= 0)
		log_error("Could not verify SSL CA");

	/* verify private key */
    if (!SSL_CTX_check_private_key(ctx))
        log_error("Check Private Key failed");
}

void sendsourceinfo(int clientfd, int serverfd)
{	
	// Create Source Information
	union
	{
		struct sockaddr sa;
		struct sockaddr_in sa4;
	} address;
	socklen_t size = sizeof(address);
		
	if (getpeername(clientfd, &address.sa, &size) >= 0)
	{
		// Write Souce Information to the server
		char sourceinfo[SOURCEDETAILSIZE];
		memset(sourceinfo, 0, SOURCEDETAILSIZE);
		
		int len = 0;
		sourceinfo[0] = STX;
		len = len + 1;
		memcpy(sourceinfo + len, g_proxyconfig.ListeningIP, strlen(g_proxyconfig.ListeningIP));
		len = len + strlen(g_proxyconfig.ListeningIP);
		sourceinfo[len] = DEL;
		len = len + 1;
		memcpy(sourceinfo + len, g_proxyconfig.SourceDetailFormat, strlen(g_proxyconfig.SourceDetailFormat));
		len = len + strlen(g_proxyconfig.SourceDetailFormat);
		sourceinfo[len] = DEL;
		len = len + 1;
		memcpy(sourceinfo + len, (char *)inet_ntoa(address.sa4.sin_addr), strlen((char*)inet_ntoa(address.sa4.sin_addr)));
		len = len + strlen((char*)inet_ntoa(address.sa4.sin_addr));
		sourceinfo[len] = ETX;
		write(serverfd, sourceinfo, strlen(sourceinfo));	
		
		// Log the details to SysLog
		log_transport("REQ", sourceinfo);			
	}	
}

void logsslerror(char* task, SSL* ssl, int code)
{	
	if (code <= 0)
	{
		char taskerror[500] = {0};
		int sslerr = SSL_get_error(ssl, code);		
				
		// Please look at http://www.openssl.org/docs/ssl/SSL_get_error.html 
		// for detailed explanation of each of the err
		switch (sslerr)
		{
			case SSL_ERROR_ZERO_RETURN: 
			{
				sprintf(taskerror, "%s: %d-%s", task, sslerr, "SSL/TLS Error - Connection closed.");
				log_error(taskerror);
				break;
			}
			case SSL_ERROR_WANT_READ:
			{
				sprintf(taskerror, "%s: %d-%s", task, sslerr, "SSL/TLS Error - Want Read.");
				log_error(taskerror);
				break;
			}			
			case SSL_ERROR_WANT_WRITE:
			{
				sprintf(taskerror, "%s: %d-%s", task, sslerr, "SSL/TLS Error - Want Write.");
				log_error(taskerror);
				break;
			}
			case SSL_ERROR_WANT_CONNECT:
			{
				sprintf(taskerror, "%s: %d-%s", task, sslerr, "SSL/TLS Error - Want Connect.");
				log_error(taskerror);
				break;
			}
			case SSL_ERROR_WANT_ACCEPT:
			{
				sprintf(taskerror, "%s: %d-%s", task, sslerr, "SSL/TLS Error - Want Accept.");
				log_error(taskerror);
				break;
			}
			case SSL_ERROR_WANT_X509_LOOKUP:
			{
				sprintf(taskerror, "%s: %d-%s", task, sslerr, "SSL/TLS Error - Want X509 Lookup.");
				log_error(taskerror);
				break;
			}
			case SSL_ERROR_SYSCALL:
			{
				sprintf(taskerror, "%s: %d-%s,%s,%d", task, sslerr, "SSL/TLS - I/O Error Occurred.", ERR_error_string(ERR_get_error(), NULL), errno);
				log_error(taskerror);
				break;
			}
			case SSL_ERROR_SSL:
			{
				sprintf(taskerror, "%s: %d-%s,%s,%d", task, sslerr, "SSL/TLS - Failure in Library.", ERR_error_string(ERR_get_error(), NULL), errno);
				log_error(taskerror);
				break;
			}
			default: 
			{
				sprintf(taskerror, "%s: %d-%s", task, sslerr, "SSL/TLS - Unknown.");
				log_error(taskerror);
				break;
			}
		};
	}
}

void incrementworkercount()
{
	pthread_mutex_lock(&g_mutex);
	g_workerthreadcount++;
	pthread_mutex_unlock(&g_mutex);	
}

void decrementworkercount()
{
	pthread_mutex_lock(&g_mutex);
	g_workerthreadcount--;
	pthread_mutex_unlock(&g_mutex);	
}

int validateclientcount()
{
	int curCount = 0;
	curCount = g_workerthreadcount;
	
	if ((g_proxyconfig.WorkerThreadCount > 0) 
			&& (curCount >= g_proxyconfig.WorkerThreadCount))
	{		
		char currbuf[40] = {0};
		sprintf(currbuf, "Rejected. Current worker count: %d", curCount);
		log_info(currbuf);

		return -1;
	}	
	
	return 0;
}

int drainsocketbuffer(int fd, SSL* ssl)
{
	if (fd <= 0)
		return -1;

	shutdown(fd, SHUT_RDWR);
	
	// Read remaining data in the buffer
	while(1)
	{
		int bytes = 0;
		char buffer[1024];
		bzero(buffer, 1024);
		
		if (ssl != NULL)
			bytes = SSL_read(ssl, buffer, 1023);
		else 
			bytes = read(fd, buffer, 1023);
		if (bytes > 0) 
		{
			log_transport("DRN", buffer);
			continue;     
		}
		else 
			break;
	}	
	
	return 0;
}

SSL* setupssl(int clientfd, SSL_CTX* ctx)
{
	SSL* ssl = NULL;
	ssl = SSL_new(ctx);
	
	if (ssl == NULL)
	{
		log_error("Could not create SSL from context");
		return NULL;
	}
	
	return ssl;
}

/**** OPENSSL callbacks ****/
void lock_callback(int mode, int type, char *file, int line)
{
	if (mode & CRYPTO_LOCK) 
		pthread_mutex_lock(&g_lockarray[type]);
	else 
		pthread_mutex_unlock(&g_lockarray[type]);
}
 
unsigned long thread_id(void)
{
	unsigned long ret;
	ret = (unsigned long)pthread_self();
	return(ret);
}
 
void init_locks(void)
{
	int i;

	g_lockarray = (pthread_mutex_t *)OPENSSL_malloc(CRYPTO_num_locks() *
											sizeof(pthread_mutex_t));
	for (i = 0; i < CRYPTO_num_locks(); i++) 
	{
		pthread_mutex_init(&g_lockarray[i],NULL);
	}

	CRYPTO_set_id_callback((unsigned long (*)())thread_id);
	CRYPTO_set_locking_callback((void (*)())lock_callback);
	CRYPTO_set_dynlock_create_callback(dyn_create_func);
	CRYPTO_set_dynlock_lock_callback(dyn_lock_func);
	CRYPTO_set_dynlock_destroy_callback(dyn_destroy_func);
}
 
void kill_locks(void)
{
	int i;

	CRYPTO_set_dynlock_create_callback(NULL);
	CRYPTO_set_dynlock_lock_callback(NULL);
	CRYPTO_set_dynlock_destroy_callback(NULL); 
	CRYPTO_set_locking_callback(NULL);

	for (i = 0; i < CRYPTO_num_locks(); i++)
	{
		pthread_mutex_destroy(&g_lockarray[i]);
	}

	free(g_lockarray);
}

struct CRYPTO_dynlock_value* dyn_create_func(const char* file, int line)
{
	struct CRYPTO_dynlock_value* value = NULL;
	value = (struct CRYPTO_dynlock_value*)malloc(sizeof(struct CRYPTO_dynlock_value));
	pthread_mutex_init(&value->mutex, NULL);
	return value;
}

void dyn_destroy_func(struct CRYPTO_dynlock_value* l, const char* file, int line)
{
	pthread_mutex_destroy(&l->mutex);
	free(l);
}

void dyn_lock_func(int mode, struct CRYPTO_dynlock_value* l, const char* file, int line)
{
	if (mode & CRYPTO_LOCK)
		pthread_mutex_lock(&l->mutex);
	else 
		pthread_mutex_unlock(&l->mutex);
}
